{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "529f4422-5c3a-4bd6-abe0-a15edfc62abb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "import polars_ds as pds\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a314316",
   "metadata": {},
   "source": [
    "# This notebook illustrates the basic usage of this package\n",
    "\n",
    "You need to create an environment with this package installed to run this notebook. (usually latest version)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aef5c69-fff3-4779-9b58-f939d725f0b0",
   "metadata": {},
   "source": [
    "# Num Extensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "430fec01-5d0b-422f-b099-c86037512b6d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 13)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>f</th><th>time_idx</th><th>dummy</th><th>actual</th><th>predicted</th><th>dummy_groups</th><th>x1</th><th>x2</th><th>x3</th><th>a</th><th>b</th><th>y</th><th>y2</th></tr><tr><td>f64</td><td>i64</td><td>str</td><td>i32</td><td>f64</td><td>str</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.0</td><td>0</td><td>&quot;a&quot;</td><td>0</td><td>0.091583</td><td>&quot;a&quot;</td><td>0.027516</td><td>0.025068</td><td>0.073583</td><td>0.610364</td><td>0.579474</td><td>-0.098684</td><td>0.007555</td></tr><tr><td>0.841471</td><td>1</td><td>&quot;a&quot;</td><td>1</td><td>0.585465</td><td>&quot;a&quot;</td><td>0.947079</td><td>0.917548</td><td>0.539384</td><td>0.248219</td><td>0.909225</td><td>-0.391686</td><td>0.482099</td></tr><tr><td>0.909297</td><td>2</td><td>&quot;a&quot;</td><td>1</td><td>0.098363</td><td>&quot;a&quot;</td><td>0.198841</td><td>0.113598</td><td>0.866751</td><td>0.52313</td><td>0.392237</td><td>-1.236188</td><td>-0.009659</td></tr><tr><td>0.14112</td><td>3</td><td>&quot;a&quot;</td><td>1</td><td>0.03237</td><td>&quot;a&quot;</td><td>0.443498</td><td>0.208141</td><td>0.137899</td><td>0.351743</td><td>0.354237</td><td>-0.077879</td><td>0.137624</td></tr><tr><td>-0.756802</td><td>4</td><td>&quot;a&quot;</td><td>1</td><td>0.324095</td><td>&quot;a&quot;</td><td>0.804234</td><td>0.103371</td><td>0.885442</td><td>0.413473</td><td>0.870453</td><td>-1.176478</td><td>0.062609</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 13)\n",
       "┌───────────┬──────────┬───────┬────────┬───┬──────────┬──────────┬───────────┬───────────┐\n",
       "│ f         ┆ time_idx ┆ dummy ┆ actual ┆ … ┆ a        ┆ b        ┆ y         ┆ y2        │\n",
       "│ ---       ┆ ---      ┆ ---   ┆ ---    ┆   ┆ ---      ┆ ---      ┆ ---       ┆ ---       │\n",
       "│ f64       ┆ i64      ┆ str   ┆ i32    ┆   ┆ f64      ┆ f64      ┆ f64       ┆ f64       │\n",
       "╞═══════════╪══════════╪═══════╪════════╪═══╪══════════╪══════════╪═══════════╪═══════════╡\n",
       "│ 0.0       ┆ 0        ┆ a     ┆ 0      ┆ … ┆ 0.610364 ┆ 0.579474 ┆ -0.098684 ┆ 0.007555  │\n",
       "│ 0.841471  ┆ 1        ┆ a     ┆ 1      ┆ … ┆ 0.248219 ┆ 0.909225 ┆ -0.391686 ┆ 0.482099  │\n",
       "│ 0.909297  ┆ 2        ┆ a     ┆ 1      ┆ … ┆ 0.52313  ┆ 0.392237 ┆ -1.236188 ┆ -0.009659 │\n",
       "│ 0.14112   ┆ 3        ┆ a     ┆ 1      ┆ … ┆ 0.351743 ┆ 0.354237 ┆ -0.077879 ┆ 0.137624  │\n",
       "│ -0.756802 ┆ 4        ┆ a     ┆ 1      ┆ … ┆ 0.413473 ┆ 0.870453 ┆ -1.176478 ┆ 0.062609  │\n",
       "└───────────┴──────────┴───────┴────────┴───┴──────────┴──────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "size = 10_000\n",
    "df = pl.DataFrame({\n",
    "    \"f\": np.sin(list(range(size)))\n",
    "    , \"time_idx\": range(size)\n",
    "    , \"dummy\": [\"a\"] * (size // 2) + [\"b\"] * (size // 2)\n",
    "    , \"actual\": np.round(np.random.random(size=size)).astype(np.int32)\n",
    "    , \"predicted\": np.random.random(size=size)\n",
    "    , \"dummy_groups\":[\"a\"] * (size//2) + [\"b\"] * (size//2) \n",
    "}).with_columns(\n",
    "    pds.random(0., 1.).alias(\"x1\")\n",
    "    , pds.random(0., 1.).alias(\"x2\")\n",
    "    , pds.random(0., 1.).alias(\"x3\")\n",
    "    , pds.random(0., 1.).alias(\"a\")\n",
    "    , pds.random(0., 1.).alias(\"b\")\n",
    ").with_columns(\n",
    "    y = pl.col(\"x1\") * 0.15 + pl.col(\"x2\") * 0.3 - pl.col(\"x3\") * 1.5 + pds.random() * 0.0001,\n",
    "    y2 = pl.col(\"x1\") * 0.13 + pl.col(\"x2\") * 0.45 - pl.col(\"x3\") * 0.1 + pds.random() * 0.0001\n",
    ")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b6f98453-34cd-4afc-b35d-db58fa60a69a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>x1</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────┐\n",
       "│ x1  │\n",
       "│ --- │\n",
       "│ f64 │\n",
       "╞═════╡\n",
       "│ 0.0 │\n",
       "└─────┘"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Column-wise Jaccard Similarity. Result should be 0 as they are distinct\n",
    "df.select(\n",
    "    pds.jaccard_col(\"x1\", pl.col(\"x2\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "416d5346-e75b-4769-a953-e898d6a4d84c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>f</th></tr><tr><td>array[f64, 2]</td></tr></thead><tbody><tr><td>[1.939505, 0.0]</td></tr><tr><td>[1.939506, 0.000209]</td></tr><tr><td>[1.939508, 0.000418]</td></tr><tr><td>[1.939512, 0.000627]</td></tr><tr><td>[1.939518, 0.000835]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────────────────┐\n",
       "│ f                    │\n",
       "│ ---                  │\n",
       "│ array[f64, 2]        │\n",
       "╞══════════════════════╡\n",
       "│ [1.939505, 0.0]      │\n",
       "│ [1.939506, 0.000209] │\n",
       "│ [1.939508, 0.000418] │\n",
       "│ [1.939512, 0.000627] │\n",
       "│ [1.939518, 0.000835] │\n",
       "└──────────────────────┘"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# FFT. First is real part, second is complex part\n",
    "# By default, this behaves the same as np's rfft, which returns a non-redundant \n",
    "# compact representation of fft output.\n",
    "df.select(\n",
    "    pds.rfft(\"f\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "71c76353",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(10000, 1)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# FFT. But return the full length\n",
    "df.select(\n",
    "    pds.rfft(\"f\", return_full=True)\n",
    ").shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cd6662d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>f</th><th>a</th><th>b</th></tr><tr><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>1.3944e-15</td><td>-0.610364</td><td>-0.579474</td></tr><tr><td>-0.841471</td><td>-0.248219</td><td>-0.909225</td></tr><tr><td>-0.909297</td><td>-0.52313</td><td>-0.392237</td></tr><tr><td>-0.14112</td><td>-0.351743</td><td>-0.354237</td></tr><tr><td>0.756802</td><td>0.196891</td><td>-0.290979</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌────────────┬───────────┬───────────┐\n",
       "│ f          ┆ a         ┆ b         │\n",
       "│ ---        ┆ ---       ┆ ---       │\n",
       "│ f64        ┆ f64       ┆ f64       │\n",
       "╞════════════╪═══════════╪═══════════╡\n",
       "│ 1.3944e-15 ┆ -0.610364 ┆ -0.579474 │\n",
       "│ -0.841471  ┆ -0.248219 ┆ -0.909225 │\n",
       "│ -0.909297  ┆ -0.52313  ┆ -0.392237 │\n",
       "│ -0.14112   ┆ -0.351743 ┆ -0.354237 │\n",
       "│ 0.756802   ┆ 0.196891  ┆ -0.290979 │\n",
       "└────────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Multiple Convolutions at once\n",
    "# Modes: `same`, `left` (left-aligned same), `right` (right-aligned same), `valid` or `full`\n",
    "# Method: `fft`, `direct`\n",
    "# Currently slower than SciPy but provides parallelism because of Polars\n",
    "df.select(\n",
    "    pds.convolve(\"f\", [-1, 0, 0, 0, 1], mode = \"full\", method = \"fft\"), # column f with the kernel given here\n",
    "    pds.convolve(\"a\", [-1, 0, 0, 0, 1], mode = \"full\", method = \"direct\"),\n",
    "    pds.convolve(\"b\", [-1, 0, 0, 0, 1], mode = \"full\", method = \"direct\"),\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ed47b643-6bcc-43f6-9a25-82168c33e7fd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>coeffs</th></tr><tr><td>list[f64]</td></tr></thead><tbody><tr><td>[-0.498886, -0.35376]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌───────────────────────┐\n",
       "│ coeffs                │\n",
       "│ ---                   │\n",
       "│ list[f64]             │\n",
       "╞═══════════════════════╡\n",
       "│ [-0.498886, -0.35376] │\n",
       "└───────────────────────┘"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Linear Regression\n",
    "df.select(\n",
    "    pds.lin_reg(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "e94324b7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>target_0</th><th>target_1</th></tr><tr><td>list[f64]</td><td>list[f64]</td></tr></thead><tbody><tr><td>[-0.498886, -0.35376]</td><td>[0.086782, 0.406454]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 2)\n",
       "┌───────────────────────┬──────────────────────┐\n",
       "│ target_0              ┆ target_1             │\n",
       "│ ---                   ┆ ---                  │\n",
       "│ list[f64]             ┆ list[f64]            │\n",
       "╞═══════════════════════╪══════════════════════╡\n",
       "│ [-0.498886, -0.35376] ┆ [0.086782, 0.406454] │\n",
       "└───────────────────────┴──────────────────────┘"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Linear Regression, multi-target\n",
    "df.select(\n",
    "    pds.lin_reg(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = [pl.col(\"y\"), pl.col(\"y2\")],\n",
    "        add_bias=False\n",
    "    )\n",
    ").unnest(\"coeffs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "7e6da23d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (4, 7)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>features</th><th>beta</th><th>std_err</th><th>t</th><th>p&gt;|t|</th><th>0.025</th><th>0.975</th></tr><tr><td>str</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>&quot;ln(x1+1)&quot;</td><td>0.219035</td><td>0.001699</td><td>128.887987</td><td>0.0</td><td>0.215704</td><td>0.222366</td></tr><tr><td>&quot;exp(x2)&quot;</td><td>0.173641</td><td>0.000686</td><td>253.137788</td><td>0.0</td><td>0.172296</td><td>0.174986</td></tr><tr><td>&quot;sin(x3)&quot;</td><td>-1.743404</td><td>0.001351</td><td>-1290.252947</td><td>0.0</td><td>-1.746052</td><td>-1.740755</td></tr><tr><td>&quot;__bias__&quot;</td><td>-0.106227</td><td>0.001517</td><td>-70.020282</td><td>0.0</td><td>-0.109201</td><td>-0.103253</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (4, 7)\n",
       "┌──────────┬───────────┬──────────┬──────────────┬───────┬───────────┬───────────┐\n",
       "│ features ┆ beta      ┆ std_err  ┆ t            ┆ p>|t| ┆ 0.025     ┆ 0.975     │\n",
       "│ ---      ┆ ---       ┆ ---      ┆ ---          ┆ ---   ┆ ---       ┆ ---       │\n",
       "│ str      ┆ f64       ┆ f64      ┆ f64          ┆ f64   ┆ f64       ┆ f64       │\n",
       "╞══════════╪═══════════╪══════════╪══════════════╪═══════╪═══════════╪═══════════╡\n",
       "│ ln(x1+1) ┆ 0.219035  ┆ 0.001699 ┆ 128.887987   ┆ 0.0   ┆ 0.215704  ┆ 0.222366  │\n",
       "│ exp(x2)  ┆ 0.173641  ┆ 0.000686 ┆ 253.137788   ┆ 0.0   ┆ 0.172296  ┆ 0.174986  │\n",
       "│ sin(x3)  ┆ -1.743404 ┆ 0.001351 ┆ -1290.252947 ┆ 0.0   ┆ -1.746052 ┆ -1.740755 │\n",
       "│ __bias__ ┆ -0.106227 ┆ 0.001517 ┆ -70.020282   ┆ 0.0   ┆ -0.109201 ┆ -0.103253 │\n",
       "└──────────┴───────────┴──────────┴──────────────┴───────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    pds.lin_reg_report(\n",
    "        # formulaic input is also available for lstsq related queries, \n",
    "        # or you can always use polars expressions, e.g. pl.col('x1') + 1, pl.col('x2').exp(), pl.col('x3').sin()\n",
    "        \"ln(x1+1)\", \"exp(x2)\", \"sin(x3)\",\n",
    "        target = \"y\",\n",
    "        add_bias = True\n",
    "    ).alias(\"report\")\n",
    ").unnest(\"report\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "16511624-fc7f-45fc-b28e-ad9a91c1bfe9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (10_000, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy</th><th>coeffs</th></tr><tr><td>str</td><td>list[f64]</td></tr></thead><tbody><tr><td>&quot;a&quot;</td><td>[-0.508478, -0.340153]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.508478, -0.340153]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.508478, -0.340153]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.508478, -0.340153]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.508478, -0.340153]</td></tr><tr><td>&hellip;</td><td>&hellip;</td></tr><tr><td>&quot;b&quot;</td><td>[-0.489381, -0.36711]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.489381, -0.36711]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.489381, -0.36711]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.489381, -0.36711]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.489381, -0.36711]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (10_000, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ coeffs                 │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ a     ┆ [-0.508478, -0.340153] │\n",
       "│ a     ┆ [-0.508478, -0.340153] │\n",
       "│ a     ┆ [-0.508478, -0.340153] │\n",
       "│ a     ┆ [-0.508478, -0.340153] │\n",
       "│ a     ┆ [-0.508478, -0.340153] │\n",
       "│ …     ┆ …                      │\n",
       "│ b     ┆ [-0.489381, -0.36711]  │\n",
       "│ b     ┆ [-0.489381, -0.36711]  │\n",
       "│ b     ┆ [-0.489381, -0.36711]  │\n",
       "│ b     ┆ [-0.489381, -0.36711]  │\n",
       "│ b     ┆ [-0.489381, -0.36711]  │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    \"dummy\",\n",
    "    pds.lin_reg(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    ).over(pl.col(\"dummy\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f550c7c6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>x1</th><th>x2</th><th>y</th><th>pred</th><th>resid</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.027516</td><td>0.025068</td><td>-0.098684</td><td>-0.022595</td><td>-0.076089</td></tr><tr><td>0.947079</td><td>0.917548</td><td>-0.391686</td><td>-0.797076</td><td>0.405389</td></tr><tr><td>0.198841</td><td>0.113598</td><td>-1.236188</td><td>-0.139386</td><td>-1.096802</td></tr><tr><td>0.443498</td><td>0.208141</td><td>-0.077879</td><td>-0.294887</td><td>0.217008</td></tr><tr><td>0.804234</td><td>0.103371</td><td>-1.176478</td><td>-0.437789</td><td>-0.738689</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 5)\n",
       "┌──────────┬──────────┬───────────┬───────────┬───────────┐\n",
       "│ x1       ┆ x2       ┆ y         ┆ pred      ┆ resid     │\n",
       "│ ---      ┆ ---      ┆ ---       ┆ ---       ┆ ---       │\n",
       "│ f64      ┆ f64      ┆ f64       ┆ f64       ┆ f64       │\n",
       "╞══════════╪══════════╪═══════════╪═══════════╪═══════════╡\n",
       "│ 0.027516 ┆ 0.025068 ┆ -0.098684 ┆ -0.022595 ┆ -0.076089 │\n",
       "│ 0.947079 ┆ 0.917548 ┆ -0.391686 ┆ -0.797076 ┆ 0.405389  │\n",
       "│ 0.198841 ┆ 0.113598 ┆ -1.236188 ┆ -0.139386 ┆ -1.096802 │\n",
       "│ 0.443498 ┆ 0.208141 ┆ -0.077879 ┆ -0.294887 ┆ 0.217008  │\n",
       "│ 0.804234 ┆ 0.103371 ┆ -1.176478 ┆ -0.437789 ┆ -0.738689 │\n",
       "└──────────┴──────────┴───────────┴───────────┴───────────┘"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# If you want prediction and residue instead of coefficients\n",
    "df.select(\n",
    "    \"x1\",\n",
    "    \"x2\",\n",
    "    \"y\",\n",
    "    pds.lin_reg(\n",
    "        \"x1\", pl.col(\"x2\"),\n",
    "        target = \"y\",\n",
    "        add_bias=False, \n",
    "        return_pred=True\n",
    "    ).alias(\"prediction\")\n",
    ").unnest(\"prediction\").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0e9fb061-340d-423d-9107-772387006ff2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy</th><th>coeffs</th></tr><tr><td>str</td><td>list[f64]</td></tr></thead><tbody><tr><td>&quot;b&quot;</td><td>[-0.489381, -0.36711]</td></tr><tr><td>&quot;a&quot;</td><td>[-0.508478, -0.340153]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ coeffs                 │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ b     ┆ [-0.489381, -0.36711]  │\n",
       "│ a     ┆ [-0.508478, -0.340153] │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.group_by(\"dummy\").agg(\n",
    "    pds.lin_reg(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        add_bias=False\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "955f2db2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy</th><th>coeffs</th></tr><tr><td>str</td><td>list[f64]</td></tr></thead><tbody><tr><td>&quot;a&quot;</td><td>[-0.343272, -0.157735]</td></tr><tr><td>&quot;b&quot;</td><td>[-0.315854, -0.193592]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬────────────────────────┐\n",
       "│ dummy ┆ coeffs                 │\n",
       "│ ---   ┆ ---                    │\n",
       "│ str   ┆ list[f64]              │\n",
       "╞═══════╪════════════════════════╡\n",
       "│ a     ┆ [-0.343272, -0.157735] │\n",
       "│ b     ┆ [-0.315854, -0.193592] │\n",
       "└───────┴────────────────────────┘"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Lasso\n",
    "df.group_by(\"dummy\").agg(\n",
    "    pds.lin_reg(\n",
    "        pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "        target = pl.col(\"y\"),\n",
    "        l1_reg = 0.1,\n",
    "        add_bias=False\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1bdae8e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy</th><th>lasso_r2</th></tr><tr><td>str</td><td>f64</td></tr></thead><tbody><tr><td>&quot;a&quot;</td><td>-0.54074</td></tr><tr><td>&quot;b&quot;</td><td>-0.548295</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌───────┬───────────┐\n",
       "│ dummy ┆ lasso_r2  │\n",
       "│ ---   ┆ ---       │\n",
       "│ str   ┆ f64       │\n",
       "╞═══════╪═══════════╡\n",
       "│ a     ┆ -0.54074  │\n",
       "│ b     ┆ -0.548295 │\n",
       "└───────┴───────────┘"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# R2 metric of lasso regressions on each group\n",
    "df.group_by(\"dummy\").agg(\n",
    "    pds.query_r2(\n",
    "        actual = pl.col(\"y\"),\n",
    "        pred = pds.lin_reg(\n",
    "            pl.col(\"x1\"), pl.col(\"x2\"),\n",
    "            target = pl.col(\"y\"),\n",
    "            l1_reg = 0.1,\n",
    "            return_pred = True,\n",
    "            add_bias=False\n",
    "        ).struct.field(\"pred\")\n",
    "    ).alias(\"lasso_r2\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "765ff27b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (10_000, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>y</th><th>x1</th><th>x2</th><th>coeffs</th><th>pred</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>list[f64]</td><td>f64</td></tr></thead><tbody><tr><td>-0.098684</td><td>0.027516</td><td>0.025068</td><td>null</td><td>null</td></tr><tr><td>-0.391686</td><td>0.947079</td><td>0.917548</td><td>null</td><td>null</td></tr><tr><td>-1.236188</td><td>0.198841</td><td>0.113598</td><td>null</td><td>null</td></tr><tr><td>-0.077879</td><td>0.443498</td><td>0.208141</td><td>null</td><td>null</td></tr><tr><td>-1.176478</td><td>0.804234</td><td>0.103371</td><td>[-1.609748, 1.186046]</td><td>-1.172012</td></tr><tr><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td><td>&hellip;</td></tr><tr><td>0.237618</td><td>0.71738</td><td>0.729978</td><td>[-0.418405, -0.473687]</td><td>-0.645937</td></tr><tr><td>-0.879749</td><td>0.388987</td><td>0.291635</td><td>[-0.813367, -0.190164]</td><td>-0.371848</td></tr><tr><td>-0.302075</td><td>0.00809</td><td>0.953496</td><td>[-0.891931, -0.00105]</td><td>-0.008217</td></tr><tr><td>-1.037887</td><td>0.229935</td><td>0.373374</td><td>[-1.01028, -0.033456]</td><td>-0.244791</td></tr><tr><td>0.163498</td><td>0.6866</td><td>0.724015</td><td>[0.116038, -0.35731]</td><td>-0.179026</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (10_000, 5)\n",
       "┌───────────┬──────────┬──────────┬────────────────────────┬───────────┐\n",
       "│ y         ┆ x1       ┆ x2       ┆ coeffs                 ┆ pred      │\n",
       "│ ---       ┆ ---      ┆ ---      ┆ ---                    ┆ ---       │\n",
       "│ f64       ┆ f64      ┆ f64      ┆ list[f64]              ┆ f64       │\n",
       "╞═══════════╪══════════╪══════════╪════════════════════════╪═══════════╡\n",
       "│ -0.098684 ┆ 0.027516 ┆ 0.025068 ┆ null                   ┆ null      │\n",
       "│ -0.391686 ┆ 0.947079 ┆ 0.917548 ┆ null                   ┆ null      │\n",
       "│ -1.236188 ┆ 0.198841 ┆ 0.113598 ┆ null                   ┆ null      │\n",
       "│ -0.077879 ┆ 0.443498 ┆ 0.208141 ┆ null                   ┆ null      │\n",
       "│ -1.176478 ┆ 0.804234 ┆ 0.103371 ┆ [-1.609748, 1.186046]  ┆ -1.172012 │\n",
       "│ …         ┆ …        ┆ …        ┆ …                      ┆ …         │\n",
       "│ 0.237618  ┆ 0.71738  ┆ 0.729978 ┆ [-0.418405, -0.473687] ┆ -0.645937 │\n",
       "│ -0.879749 ┆ 0.388987 ┆ 0.291635 ┆ [-0.813367, -0.190164] ┆ -0.371848 │\n",
       "│ -0.302075 ┆ 0.00809  ┆ 0.953496 ┆ [-0.891931, -0.00105]  ┆ -0.008217 │\n",
       "│ -1.037887 ┆ 0.229935 ┆ 0.373374 ┆ [-1.01028, -0.033456]  ┆ -0.244791 │\n",
       "│ 0.163498  ┆ 0.6866   ┆ 0.724015 ┆ [0.116038, -0.35731]   ┆ -0.179026 │\n",
       "└───────────┴──────────┴──────────┴────────────────────────┴───────────┘"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Rolling regression\n",
    "df.select(\n",
    "    \"y\",\n",
    "    \"x1\",\n",
    "    \"x2\",\n",
    "    pds.rolling_lin_reg(\n",
    "        \"x1\", \"x2\",\n",
    "        target = \"y\",\n",
    "        window_size = 5,\n",
    "        null_policy = \"zero\"\n",
    "    ).alias(\"result\")\n",
    ").unnest(\"result\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d8fda8ca-57e7-4e02-a3f0-283ecce66a59",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>y</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>-0.0</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌──────┐\n",
       "│ y    │\n",
       "│ ---  │\n",
       "│ f64  │\n",
       "╞══════╡\n",
       "│ -0.0 │\n",
       "└──────┘"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Conditional Entropy, should be 0 because x1 is an ID\n",
    "df.select(\n",
    "    pds.query_cond_entropy(\"y\", \"x1\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "81def1cf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th></tr><tr><td>list[f64]</td></tr></thead><tbody><tr><td>[28.850744, 28.801703, 28.618474]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌─────────────────────────────────┐\n",
       "│ a                               │\n",
       "│ ---                             │\n",
       "│ list[f64]                       │\n",
       "╞═════════════════════════════════╡\n",
       "│ [28.850744, 28.801703, 28.6184… │\n",
       "└─────────────────────────────────┘"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Only want singular values (principal values?)\n",
    "df.select(\n",
    "    pds.singular_values(\"a\", \"b\", \"x1\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "cc497383",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>singular_value</th><th>weight_vector</th></tr><tr><td>f64</td><td>list[f64]</td></tr></thead><tbody><tr><td>28.820497</td><td>[0.995171, 0.098156]</td></tr><tr><td>28.76974</td><td>[-0.098156, 0.995171]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌────────────────┬───────────────────────┐\n",
       "│ singular_value ┆ weight_vector         │\n",
       "│ ---            ┆ ---                   │\n",
       "│ f64            ┆ list[f64]             │\n",
       "╞════════════════╪═══════════════════════╡\n",
       "│ 28.820497      ┆ [0.995171, 0.098156]  │\n",
       "│ 28.76974       ┆ [-0.098156, 0.995171] │\n",
       "└────────────────┴───────────────────────┘"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Singular values + The principal components\n",
    "df.select(\n",
    "    pds.pca(\"a\", \"b\")\n",
    ").unnest(\"a\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e047d40a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>pc1</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.11709</td></tr><tr><td>-0.210939</td></tr><tr><td>0.011899</td></tr><tr><td>-0.162391</td></tr><tr><td>-0.050289</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌───────────┐\n",
       "│ pc1       │\n",
       "│ ---       │\n",
       "│ f64       │\n",
       "╞═══════════╡\n",
       "│ 0.11709   │\n",
       "│ -0.210939 │\n",
       "│ 0.011899  │\n",
       "│ -0.162391 │\n",
       "│ -0.050289 │\n",
       "└───────────┘"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# PC1\n",
    "df.select(\n",
    "    pds.principal_components(\"a\", \"b\", k =1).alias(\"principal_components\")\n",
    ").unnest(\"principal_components\").head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b2e036f",
   "metadata": {},
   "source": [
    "# ML Metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "85d0d094-3c4c-4230-a589-1027c5690162",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 8)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>dummy_groups</th><th>l2</th><th>log loss</th><th>precision</th><th>recall</th><th>f</th><th>average_precision</th><th>roc_auc</th></tr><tr><td>str</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>&quot;b&quot;</td><td>0.334676</td><td>1.002754</td><td>0.490268</td><td>0.484181</td><td>0.487205</td><td>0.503104</td><td>0.500475</td></tr><tr><td>&quot;a&quot;</td><td>0.328401</td><td>0.985268</td><td>0.508709</td><td>0.512565</td><td>0.51063</td><td>0.504645</td><td>0.507109</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 8)\n",
       "┌──────────────┬──────────┬──────────┬───────────┬──────────┬──────────┬────────────────┬──────────┐\n",
       "│ dummy_groups ┆ l2       ┆ log loss ┆ precision ┆ recall   ┆ f        ┆ average_precis ┆ roc_auc  │\n",
       "│ ---          ┆ ---      ┆ ---      ┆ ---       ┆ ---      ┆ ---      ┆ ion            ┆ ---      │\n",
       "│ str          ┆ f64      ┆ f64      ┆ f64       ┆ f64      ┆ f64      ┆ ---            ┆ f64      │\n",
       "│              ┆          ┆          ┆           ┆          ┆          ┆ f64            ┆          │\n",
       "╞══════════════╪══════════╪══════════╪═══════════╪══════════╪══════════╪════════════════╪══════════╡\n",
       "│ b            ┆ 0.334676 ┆ 1.002754 ┆ 0.490268  ┆ 0.484181 ┆ 0.487205 ┆ 0.503104       ┆ 0.500475 │\n",
       "│ a            ┆ 0.328401 ┆ 0.985268 ┆ 0.508709  ┆ 0.512565 ┆ 0.51063  ┆ 0.504645       ┆ 0.507109 │\n",
       "└──────────────┴──────────┴──────────┴───────────┴──────────┴──────────┴────────────────┴──────────┘"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.group_by(\"dummy_groups\").agg(\n",
    "    pds.query_l2(\"actual\", \"predicted\").alias(\"l2\"),\n",
    "    pds.query_log_loss(\"actual\", \"predicted\").alias(\"log loss\"),\n",
    "    pds.query_binary_metrics(actual=\"actual\", pred=\"predicted\").alias(\"combo\")\n",
    ").unnest(\"combo\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8d7c6e3-0f1d-45f0-9fdb-cdb303b98556",
   "metadata": {},
   "source": [
    "# Str Extension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "54ad36f9-264e-4a49-bf36-936639440edf",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>sen</th><th>word</th></tr><tr><td>str</td><td>str</td></tr></thead><tbody><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;words&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;words&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;words&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌─────────────────────────────────┬───────┐\n",
       "│ sen                             ┆ word  │\n",
       "│ ---                             ┆ ---   │\n",
       "│ str                             ┆ str   │\n",
       "╞═════════════════════════════════╪═══════╡\n",
       "│ Hello, world! I'm going to chu… ┆ words │\n",
       "│ Hello, world! I'm going to chu… ┆ word  │\n",
       "│ Hello, world! I'm going to chu… ┆ words │\n",
       "│ Hello, world! I'm going to chu… ┆ word  │\n",
       "│ Hello, world! I'm going to chu… ┆ words │\n",
       "└─────────────────────────────────┴───────┘"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "size = 100_000\n",
    "df2 = pl.DataFrame({\n",
    "    \"sen\":[\"Hello, world! I'm going to church.\"] * size,\n",
    "    \"word\":[\"words\", \"word\"] * (size //2)\n",
    "})\n",
    "df2.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ee123a7e-7f9b-4f48-a5d5-6354799201ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>sen</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;hello&quot;</td></tr><tr><td>&quot;going&quot;</td></tr><tr><td>&quot;world&quot;</td></tr><tr><td>&quot;church&quot;</td></tr><tr><td>&quot;to&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌────────┐\n",
       "│ sen    │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│ hello  │\n",
       "│ going  │\n",
       "│ world  │\n",
       "│ church │\n",
       "│ to     │\n",
       "└────────┘"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Tokenize\n",
    "df2.select(\n",
    "    pds.str_tokenize(pl.col(\"sen\").str.to_lowercase()).explode().unique()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f33017e3-17df-498b-93d9-1d656a344388",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>sen</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;go&quot;</td></tr><tr><td>&quot;hello&quot;</td></tr><tr><td>&quot;world&quot;</td></tr><tr><td>&quot;church&quot;</td></tr><tr><td>&quot;&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌────────┐\n",
       "│ sen    │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│ go     │\n",
       "│ hello  │\n",
       "│ world  │\n",
       "│ church │\n",
       "│        │\n",
       "└────────┘"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.select(\n",
    "    pds.str_tokenize(pl.col(\"sen\").str.to_lowercase(), stem=True).explode().unique()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "69237c02-5f9f-4e92-b68d-6ac43aad1a79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th></tr><tr><td>u32</td></tr></thead><tbody><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────┐\n",
       "│ word │\n",
       "│ ---  │\n",
       "│ u32  │\n",
       "╞══════╡\n",
       "│ 2    │\n",
       "│ 1    │\n",
       "│ 2    │\n",
       "│ 1    │\n",
       "│ 2    │\n",
       "└──────┘"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.select(\n",
    "    pds.str_leven(\"word\", pl.lit(\"world\"))\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "2eba320c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th></tr><tr><td>u32</td></tr></thead><tbody><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────┐\n",
       "│ word │\n",
       "│ ---  │\n",
       "│ u32  │\n",
       "╞══════╡\n",
       "│ 2    │\n",
       "│ 1    │\n",
       "│ 2    │\n",
       "│ 1    │\n",
       "│ 2    │\n",
       "└──────┘"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Damerau-Levenshtein\n",
    "df2.select(\n",
    "    pds.str_d_leven(\"word\", pl.lit(\"world\"))\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "795396dc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.6</td></tr><tr><td>0.8</td></tr><tr><td>0.6</td></tr><tr><td>0.8</td></tr><tr><td>0.6</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────┐\n",
       "│ word │\n",
       "│ ---  │\n",
       "│ f64  │\n",
       "╞══════╡\n",
       "│ 0.6  │\n",
       "│ 0.8  │\n",
       "│ 0.6  │\n",
       "│ 0.8  │\n",
       "│ 0.6  │\n",
       "└──────┘"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.select( # column \"word\" vs. the word \"world\"\n",
    "    pds.str_leven(\"word\", pl.lit(\"world\"), return_sim = True)\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "2dad7633-67fa-47f3-b86a-9f4cd097a650",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>sen</th><th>word</th></tr><tr><td>str</td><td>str</td></tr></thead><tbody><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr><tr><td>&quot;Hello, world! I&#x27;m going to chu…</td><td>&quot;word&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌─────────────────────────────────┬──────┐\n",
       "│ sen                             ┆ word │\n",
       "│ ---                             ┆ ---  │\n",
       "│ str                             ┆ str  │\n",
       "╞═════════════════════════════════╪══════╡\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "│ Hello, world! I'm going to chu… ┆ word │\n",
       "└─────────────────────────────────┴──────┘"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df2.filter(\n",
    "    # This is way faster than computing ditance and then doing a filter\n",
    "    pds.filter_by_levenshtein(pl.col(\"word\"), pl.lit(\"world\"), 1) # <= 1. \n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "dc9477c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.DataFrame({\n",
    "    \"word\":[\"apple\", \"banana\", \"pineapple\", \"asasasas\", \"sasasass\"],\n",
    "    \"other_data\": [1,2,3,4,5]\n",
    "})\n",
    "gibberish = [\"asasasa\", \"sasaaasss\", \"asdasadadfa\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "cf0c0e72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th></tr><tr><td>str</td></tr></thead><tbody><tr><td>&quot;banana&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 1)\n",
       "┌────────┐\n",
       "│ word   │\n",
       "│ ---    │\n",
       "│ str    │\n",
       "╞════════╡\n",
       "│ banana │\n",
       "└────────┘"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    # Nearest string\n",
    "    pds.str_nearest(\"word\", word = \"banana\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "c50591e0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (2, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>word</th><th>other_data</th></tr><tr><td>str</td><td>i64</td></tr></thead><tbody><tr><td>&quot;asasasas&quot;</td><td>4</td></tr><tr><td>&quot;sasasass&quot;</td><td>5</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (2, 2)\n",
       "┌──────────┬────────────┐\n",
       "│ word     ┆ other_data │\n",
       "│ ---      ┆ ---        │\n",
       "│ str      ┆ i64        │\n",
       "╞══════════╪════════════╡\n",
       "│ asasasas ┆ 4          │\n",
       "│ sasasass ┆ 5          │\n",
       "└──────────┴────────────┘"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.filter(\n",
    "    # Filters to words that are similar to any word in vocab\n",
    "    pds.similar_to_vocab(\n",
    "        pl.col(\"word\"),\n",
    "        vocab = gibberish,\n",
    "        threshold = 0.5,\n",
    "        metric = \"lv\", # Levenshtein similarity. Other options: dleven, osa, jw\n",
    "        strategy = \"any\" # True if the word is similar to any word in vocab. Other options: \"all\", \"avg\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "7ece3794",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>asasasa</th><th>sasaaasss</th><th>asdasadadfa</th><th>LCS based Fuzz match - apples</th><th>Optimal String Alignment - apples</th><th>Jaro-Winkler - apples</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.142857</td><td>0.111111</td><td>0.090909</td><td>0.833333</td><td>0.833333</td><td>0.966667</td></tr><tr><td>0.428571</td><td>0.333333</td><td>0.272727</td><td>0.166667</td><td>0.0</td><td>0.444444</td></tr><tr><td>0.111111</td><td>0.111111</td><td>0.090909</td><td>0.555556</td><td>0.444444</td><td>0.5</td></tr><tr><td>0.875</td><td>0.666667</td><td>0.545455</td><td>0.25</td><td>0.25</td><td>0.527778</td></tr><tr><td>0.75</td><td>0.777778</td><td>0.454545</td><td>0.25</td><td>0.25</td><td>0.527778</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌──────────┬───────────┬─────────────┬────────────────┬──────────────────┬─────────────────────────┐\n",
       "│ asasasa  ┆ sasaaasss ┆ asdasadadfa ┆ LCS based Fuzz ┆ Optimal String   ┆ Jaro-Winkler - apples   │\n",
       "│ ---      ┆ ---       ┆ ---         ┆ match - apples ┆ Alignment - app… ┆ ---                     │\n",
       "│ f64      ┆ f64       ┆ f64         ┆ ---            ┆ ---              ┆ f64                     │\n",
       "│          ┆           ┆             ┆ f64            ┆ f64              ┆                         │\n",
       "╞══════════╪═══════════╪═════════════╪════════════════╪══════════════════╪═════════════════════════╡\n",
       "│ 0.142857 ┆ 0.111111  ┆ 0.090909    ┆ 0.833333       ┆ 0.833333         ┆ 0.966667                │\n",
       "│ 0.428571 ┆ 0.333333  ┆ 0.272727    ┆ 0.166667       ┆ 0.0              ┆ 0.444444                │\n",
       "│ 0.111111 ┆ 0.111111  ┆ 0.090909    ┆ 0.555556       ┆ 0.444444         ┆ 0.5                     │\n",
       "│ 0.875    ┆ 0.666667  ┆ 0.545455    ┆ 0.25           ┆ 0.25             ┆ 0.527778                │\n",
       "│ 0.75     ┆ 0.777778  ┆ 0.454545    ┆ 0.25           ┆ 0.25             ┆ 0.527778                │\n",
       "└──────────┴───────────┴─────────────┴────────────────┴──────────────────┴─────────────────────────┘"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.select(\n",
    "    pds.str_leven(\"word\", pl.lit(\"asasasa\"), return_sim=True).alias(\"asasasa\"),\n",
    "    pds.str_leven(\"word\", pl.lit(\"sasaaasss\"), return_sim=True).alias(\"sasaaasss\"),\n",
    "    pds.str_leven(\"word\", pl.lit(\"asdasadadfa\"), return_sim=True).alias(\"asdasadadfa\"),\n",
    "    pds.str_fuzz(\"word\", pl.lit(\"apples\")).alias(\"LCS based Fuzz match - apples\"),\n",
    "    pds.str_osa(\"word\", pl.lit(\"apples\"), return_sim=True).alias(\"Optimal String Alignment - apples\"),\n",
    "    pds.str_jw(\"word\", pl.lit(\"apples\")).alias(\"Jaro-Winkler - apples\"),\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8841f2a1",
   "metadata": {},
   "source": [
    "# Stats Extension"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "2c6171b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 1)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>null</td></tr><tr><td>null</td></tr><tr><td>0.166248</td></tr><tr><td>1.339555</td></tr><tr><td>1.29705</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 1)\n",
       "┌──────────┐\n",
       "│ a        │\n",
       "│ ---      │\n",
       "│ f64      │\n",
       "╞══════════╡\n",
       "│ null     │\n",
       "│ null     │\n",
       "│ 0.166248 │\n",
       "│ 1.339555 │\n",
       "│ 1.29705  │\n",
       "└──────────┘"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "df = pl.DataFrame({\n",
    "    \"a\": [None, None] + list(np.random.normal(size = 998))\n",
    "})\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "2f6e7445",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>random_normal</th><th>random_normal_that_respects_null_of_a</th></tr><tr><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>null</td><td>-0.819894</td><td>null</td></tr><tr><td>null</td><td>-0.106444</td><td>null</td></tr><tr><td>0.166248</td><td>0.311177</td><td>0.46741</td></tr><tr><td>1.339555</td><td>1.507627</td><td>1.895496</td></tr><tr><td>1.29705</td><td>-0.40433</td><td>0.528693</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌──────────┬───────────────┬─────────────────────────────────┐\n",
       "│ a        ┆ random_normal ┆ random_normal_that_respects_nu… │\n",
       "│ ---      ┆ ---           ┆ ---                             │\n",
       "│ f64      ┆ f64           ┆ f64                             │\n",
       "╞══════════╪═══════════════╪═════════════════════════════════╡\n",
       "│ null     ┆ -0.819894     ┆ null                            │\n",
       "│ null     ┆ -0.106444     ┆ null                            │\n",
       "│ 0.166248 ┆ 0.311177      ┆ 0.46741                         │\n",
       "│ 1.339555 ┆ 1.507627      ┆ 1.895496                        │\n",
       "│ 1.29705  ┆ -0.40433      ┆ 0.528693                        │\n",
       "└──────────┴───────────────┴─────────────────────────────────┘"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate random numbers, respecting null positions in reference column (pl.col(\"a\"))\n",
    "df.with_columns(\n",
    "    pds.random_normal(mean = 0.5, std = 1.0).alias(\"random_normal\"),\n",
    "    pl.when(pl.col(\"a\").is_null()).then(None).otherwise(\n",
    "        pds.random_normal(mean = 0.5, std = 1.0).alias(\"random_normal\")\n",
    "    ).alias(\"random_normal_that_respects_null_of_a\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "11e13f55",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>random_str</th><th>random_str_that_respects_null_of_a</th></tr><tr><td>f64</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>null</td><td>&quot;cC&quot;</td><td>null</td></tr><tr><td>null</td><td>&quot;HxBp&quot;</td><td>null</td></tr><tr><td>0.166248</td><td>&quot;qjl&quot;</td><td>&quot;RMFn&quot;</td></tr><tr><td>1.339555</td><td>&quot;3N&quot;</td><td>&quot;n6Al&quot;</td></tr><tr><td>1.29705</td><td>&quot;cHD&quot;</td><td>&quot;MF&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌──────────┬────────────┬─────────────────────────────────┐\n",
       "│ a        ┆ random_str ┆ random_str_that_respects_null_… │\n",
       "│ ---      ┆ ---        ┆ ---                             │\n",
       "│ f64      ┆ str        ┆ str                             │\n",
       "╞══════════╪════════════╪═════════════════════════════════╡\n",
       "│ null     ┆ cC         ┆ null                            │\n",
       "│ null     ┆ HxBp       ┆ null                            │\n",
       "│ 0.166248 ┆ qjl        ┆ RMFn                            │\n",
       "│ 1.339555 ┆ 3N         ┆ n6Al                            │\n",
       "│ 1.29705  ┆ cHD        ┆ MF                              │\n",
       "└──────────┴────────────┴─────────────────────────────────┘"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate random string\n",
    "df.with_columns(\n",
    "    pds.random_str(min_size = 1, max_size = 5).alias(\"random_str\"),\n",
    "    pl.when(pl.col(\"a\").is_null()).then(None).otherwise(\n",
    "        pds.random_str(min_size = 1, max_size = 5)\n",
    "    ).alias(\"random_str_that_respects_null_of_a\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "43c37394",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>random_str</th></tr><tr><td>f64</td><td>str</td></tr></thead><tbody><tr><td>null</td><td>null</td></tr><tr><td>null</td><td>null</td></tr><tr><td>0.166248</td><td>&quot;uhpES&quot;</td></tr><tr><td>1.339555</td><td>&quot;1AzJe&quot;</td></tr><tr><td>1.29705</td><td>&quot;EpZWF&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 2)\n",
       "┌──────────┬────────────┐\n",
       "│ a        ┆ random_str │\n",
       "│ ---      ┆ ---        │\n",
       "│ f64      ┆ str        │\n",
       "╞══════════╪════════════╡\n",
       "│ null     ┆ null       │\n",
       "│ null     ┆ null       │\n",
       "│ 0.166248 ┆ uhpES      │\n",
       "│ 1.339555 ┆ 1AzJe      │\n",
       "│ 1.29705  ┆ EpZWF      │\n",
       "└──────────┴────────────┘"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate fixed size random string, while respecting column a's nulls\n",
    "df.with_columns(\n",
    "    pl.when(pl.col(\"a\").is_null()).then(None).otherwise(\n",
    "        pds.random_str(min_size = 5, max_size = 5)\n",
    "    ).alias(\"random_str\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "3d0c06a4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>test1</th><th>literal</th><th>test1_perturbed</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>null</td><td>0.175841</td><td>null</td><td>0.176231</td></tr><tr><td>null</td><td>-0.816172</td><td>null</td><td>-0.815861</td></tr><tr><td>0.166248</td><td>1.955628</td><td>1.995267</td><td>1.955387</td></tr><tr><td>1.339555</td><td>0.576981</td><td>3.011182</td><td>0.57688</td></tr><tr><td>1.29705</td><td>-1.083462</td><td>1.005481</td><td>-1.083934</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌──────────┬───────────┬──────────┬─────────────────┐\n",
       "│ a        ┆ test1     ┆ literal  ┆ test1_perturbed │\n",
       "│ ---      ┆ ---       ┆ ---      ┆ ---             │\n",
       "│ f64      ┆ f64       ┆ f64      ┆ f64             │\n",
       "╞══════════╪═══════════╪══════════╪═════════════════╡\n",
       "│ null     ┆ 0.175841  ┆ null     ┆ 0.176231        │\n",
       "│ null     ┆ -0.816172 ┆ null     ┆ -0.815861       │\n",
       "│ 0.166248 ┆ 1.955628  ┆ 1.995267 ┆ 1.955387        │\n",
       "│ 1.339555 ┆ 0.576981  ┆ 3.011182 ┆ 0.57688         │\n",
       "│ 1.29705  ┆ -1.083462 ┆ 1.005481 ┆ -1.083934       │\n",
       "└──────────┴───────────┴──────────┴─────────────────┘"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.with_columns(\n",
    "    # Sample from a normal distribution, using reference column \"a\" 's mean and std\n",
    "    pds.random_normal(pl.col(\"a\").mean(), pl.col(\"a\").std()).alias(\"test1\") \n",
    "    # Sample from uniform distribution, with low = 0 and high = \"a\"'s max, and respect the nulls in \"a\"\n",
    "    , pl.when(pl.col(\"a\").is_null()).then(None).otherwise(\n",
    "        pds.random(lower = 0., upper = pl.col(\"a\").max()).alias(\"test2\")\n",
    "    )\n",
    ").with_columns(\n",
    "    # Add a random pertubation to test1\n",
    "    pds.perturb(\"test1\", epsilon=0.001).alias(\"test1_perturbed\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "67dc6583",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>a</th><th>[0, 1)</th><th>Normal</th><th>Int from [0, 10)</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>i32</td></tr></thead><tbody><tr><td>null</td><td>0.714802</td><td>-0.54607</td><td>2</td></tr><tr><td>null</td><td>0.355513</td><td>0.827599</td><td>9</td></tr><tr><td>0.166248</td><td>0.002282</td><td>0.570296</td><td>0</td></tr><tr><td>1.339555</td><td>0.786958</td><td>-0.031061</td><td>3</td></tr><tr><td>1.29705</td><td>0.320615</td><td>-0.158668</td><td>8</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌──────────┬──────────┬───────────┬──────────────────┐\n",
       "│ a        ┆ [0, 1)   ┆ Normal    ┆ Int from [0, 10) │\n",
       "│ ---      ┆ ---      ┆ ---       ┆ ---              │\n",
       "│ f64      ┆ f64      ┆ f64       ┆ i32              │\n",
       "╞══════════╪══════════╪═══════════╪══════════════════╡\n",
       "│ null     ┆ 0.714802 ┆ -0.54607  ┆ 2                │\n",
       "│ null     ┆ 0.355513 ┆ 0.827599  ┆ 9                │\n",
       "│ 0.166248 ┆ 0.002282 ┆ 0.570296  ┆ 0                │\n",
       "│ 1.339555 ┆ 0.786958 ┆ -0.031061 ┆ 3                │\n",
       "│ 1.29705  ┆ 0.320615 ┆ -0.158668 ┆ 8                │\n",
       "└──────────┴──────────┴───────────┴──────────────────┘"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# New in v0.3.5\n",
    "# This way, we don't have a reference column, so we cannot respect nulls, but is more convenient to use.\n",
    "df.with_columns(\n",
    "    pds.random().alias(\"[0, 1)\"),\n",
    "    pds.random_normal(pl.col(\"a\").mean(), pl.col(\"a\").std()).alias(\"Normal\"),\n",
    "    pds.random_int(0, 10).alias(\"Int from [0, 10)\"),\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "7b63f636",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>t-tests: statistics</th><th>t-tests: pvalue</th><th>normality_test: statistics</th><th>normality_test: pvalue</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>-0.425243</td><td>0.670722</td><td>0.782871</td><td>0.676086</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 4)\n",
       "┌─────────────────────┬─────────────────┬────────────────────────────┬────────────────────────┐\n",
       "│ t-tests: statistics ┆ t-tests: pvalue ┆ normality_test: statistics ┆ normality_test: pvalue │\n",
       "│ ---                 ┆ ---             ┆ ---                        ┆ ---                    │\n",
       "│ f64                 ┆ f64             ┆ f64                        ┆ f64                    │\n",
       "╞═════════════════════╪═════════════════╪════════════════════════════╪════════════════════════╡\n",
       "│ -0.425243           ┆ 0.670722        ┆ 0.782871                   ┆ 0.676086               │\n",
       "└─────────────────────┴─────────────────┴────────────────────────────┴────────────────────────┘"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Genenrate 2 random sample, both normally distributed\n",
    "# Run Welch's t test on them, p value should be big since they have equal mean\n",
    "# Run a normality test. Again, p value should be big since they are normally distributed \n",
    "\n",
    "df.with_columns(\n",
    "    pds.random_normal(0.5, 1.0).alias(\"test1\"),\n",
    "    pds.random_normal(0.5, 2.0).alias(\"test2\"),\n",
    ").select(\n",
    "    pds.ttest_ind(\"test1\", \"test2\", equal_var=False).alias(\"t-test\"),\n",
    "    pds.normal_test(\"test1\").alias(\"normality_test\")\n",
    ").select(\n",
    "    pl.col(\"t-test\").struct.field(\"statistic\").alias(\"t-tests: statistics\")\n",
    "    , pl.col(\"t-test\").struct.field(\"pvalue\").alias(\"t-tests: pvalue\")\n",
    "    , pl.col(\"normality_test\").struct.field(\"statistic\").alias(\"normality_test: statistics\")\n",
    "    , pl.col(\"normality_test\").struct.field(\"pvalue\").alias(\"normality_test: pvalue\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "b46a72a5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 5)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>market_id</th><th>var1</th><th>var2</th><th>category_1</th><th>category_2</th></tr><tr><td>i64</td><td>f64</td><td>f64</td><td>i32</td><td>i32</td></tr></thead><tbody><tr><td>0</td><td>0.598321</td><td>0.075415</td><td>4</td><td>5</td></tr><tr><td>1</td><td>0.073296</td><td>0.789893</td><td>2</td><td>3</td></tr><tr><td>2</td><td>0.818023</td><td>0.504974</td><td>2</td><td>4</td></tr><tr><td>0</td><td>0.985104</td><td>0.153053</td><td>0</td><td>6</td></tr><tr><td>1</td><td>0.440852</td><td>0.862906</td><td>2</td><td>1</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 5)\n",
       "┌───────────┬──────────┬──────────┬────────────┬────────────┐\n",
       "│ market_id ┆ var1     ┆ var2     ┆ category_1 ┆ category_2 │\n",
       "│ ---       ┆ ---      ┆ ---      ┆ ---        ┆ ---        │\n",
       "│ i64       ┆ f64      ┆ f64      ┆ i32        ┆ i32        │\n",
       "╞═══════════╪══════════╪══════════╪════════════╪════════════╡\n",
       "│ 0         ┆ 0.598321 ┆ 0.075415 ┆ 4          ┆ 5          │\n",
       "│ 1         ┆ 0.073296 ┆ 0.789893 ┆ 2          ┆ 3          │\n",
       "│ 2         ┆ 0.818023 ┆ 0.504974 ┆ 2          ┆ 4          │\n",
       "│ 0         ┆ 0.985104 ┆ 0.153053 ┆ 0          ┆ 6          │\n",
       "│ 1         ┆ 0.440852 ┆ 0.862906 ┆ 2          ┆ 1          │\n",
       "└───────────┴──────────┴──────────┴────────────┴────────────┘"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "size = 5_000\n",
    "df = pl.DataFrame({\n",
    "    \"market_id\": range(size),\n",
    "}).with_columns(\n",
    "    pl.col(\"market_id\").mod(3),\n",
    "    var1 = pds.random(),\n",
    "    var2 = pds.random(),\n",
    "    category_1 = pds.random_int(0, 5),\n",
    "    category_2 = pds.random_int(0, 10),\n",
    ")\n",
    "\n",
    "df.head(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "adc4f66f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>t-test</th><th>chi2-test</th><th>f-test</th></tr><tr><td>struct[2]</td><td>struct[2]</td><td>struct[2]</td></tr></thead><tbody><tr><td>{0.356596,0.721402}</td><td>{31.810658,0.668157}</td><td>{0.744296,0.5617}</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1, 3)\n",
       "┌─────────────────────┬──────────────────────┬───────────────────┐\n",
       "│ t-test              ┆ chi2-test            ┆ f-test            │\n",
       "│ ---                 ┆ ---                  ┆ ---               │\n",
       "│ struct[2]           ┆ struct[2]            ┆ struct[2]         │\n",
       "╞═════════════════════╪══════════════════════╪═══════════════════╡\n",
       "│ {0.356596,0.721402} ┆ {31.810658,0.668157} ┆ {0.744296,0.5617} │\n",
       "└─────────────────────┴──────────────────────┴───────────────────┘"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# In dataframe statistical tests!\n",
    "df.select(\n",
    "    pds.ttest_ind(\"var1\", \"var2\", equal_var=True).alias(\"t-test\"),\n",
    "    pds.chi2(\"category_1\", \"category_2\").alias(\"chi2-test\"),\n",
    "    pds.f_test(\"var1\", group = \"category_1\").alias(\"f-test\")\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "65dbb6bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (3, 4)\n",
      "┌───────────┬──────────────────────┬──────────────────────┬─────────────────────┐\n",
      "│ market_id ┆ t-test               ┆ chi2-test            ┆ f-test              │\n",
      "│ ---       ┆ ---                  ┆ ---                  ┆ ---                 │\n",
      "│ i64       ┆ struct[2]            ┆ struct[2]            ┆ struct[2]           │\n",
      "╞═══════════╪══════════════════════╪══════════════════════╪═════════════════════╡\n",
      "│ 0         ┆ {0.782406,0.434031}  ┆ {32.80012,0.621581}  ┆ {2.027486,0.088156} │\n",
      "│ 1         ┆ {-1.168306,0.242767} ┆ {34.251982,0.551894} ┆ {0.414089,0.798598} │\n",
      "│ 2         ┆ {0.988312,0.323072}  ┆ {35.722092,0.481702} ┆ {1.335438,0.254489} │\n",
      "└───────────┴──────────────────────┴──────────────────────┴─────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Can also be done in group by context\n",
    "print(\n",
    "    df.group_by(\"market_id\").agg(\n",
    "        pds.ttest_ind(\"var1\", \"var2\", equal_var=False).alias(\"t-test\"),\n",
    "        pds.chi2(\"category_1\", \"category_2\").alias(\"chi2-test\"),\n",
    "        pds.f_test(\"var1\", group = \"category_1\").alias(\"f-test\")\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "843d54c5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (9, 2)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>first_digit_cnt</th><th>first_digit_distribution</th></tr><tr><td>u32</td><td>f64</td></tr></thead><tbody><tr><td>537</td><td>0.1074</td></tr><tr><td>534</td><td>0.1068</td></tr><tr><td>577</td><td>0.1154</td></tr><tr><td>605</td><td>0.121</td></tr><tr><td>546</td><td>0.1092</td></tr><tr><td>572</td><td>0.1144</td></tr><tr><td>505</td><td>0.101</td></tr><tr><td>561</td><td>0.1122</td></tr><tr><td>563</td><td>0.1126</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (9, 2)\n",
       "┌─────────────────┬──────────────────────────┐\n",
       "│ first_digit_cnt ┆ first_digit_distribution │\n",
       "│ ---             ┆ ---                      │\n",
       "│ u32             ┆ f64                      │\n",
       "╞═════════════════╪══════════════════════════╡\n",
       "│ 537             ┆ 0.1074                   │\n",
       "│ 534             ┆ 0.1068                   │\n",
       "│ 577             ┆ 0.1154                   │\n",
       "│ 605             ┆ 0.121                    │\n",
       "│ 546             ┆ 0.1092                   │\n",
       "│ 572             ┆ 0.1144                   │\n",
       "│ 505             ┆ 0.101                    │\n",
       "│ 561             ┆ 0.1122                   │\n",
       "│ 563             ┆ 0.1126                   │\n",
       "└─────────────────┴──────────────────────────┘"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Benford's law\n",
    "df.select(\n",
    "    first_digit_cnt = pds.query_first_digit_cnt(pl.col(\"var1\")).explode()\n",
    ").with_columns(\n",
    "    # This doesn't follow benford's law because it is random data\n",
    "    first_digit_distribution = pl.col(\"first_digit_cnt\") / pl.col(\"first_digit_cnt\").sum()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b232e4d7",
   "metadata": {},
   "source": [
    "# Nearest Neighbors Related Tasks\n",
    "\n",
    "These queries can be very slow when data/dimension gets huge, even when processed in parallel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "87aff1ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars_ds as pds\n",
    "size = 2000\n",
    "df = pl.DataFrame({\n",
    "    \"id\": range(size), \n",
    "}).with_columns(\n",
    "    pds.random().alias(\"var1\"),\n",
    "    pds.random().alias(\"var2\"),\n",
    "    pds.random().alias(\"var3\"),\n",
    "    pds.random().alias(\"r\"),\n",
    "    (pds.random() * 10).alias(\"rh\"),\n",
    "    pl.col(\"id\").cast(pl.UInt32)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "2fae4b5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 7)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th><th>nb_l_inf_cnt</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>u32</td></tr></thead><tbody><tr><td>0</td><td>0.720043</td><td>0.762057</td><td>0.0802</td><td>0.853373</td><td>1.184888</td><td>16</td></tr><tr><td>1</td><td>0.746859</td><td>0.774783</td><td>0.969885</td><td>0.027992</td><td>6.011372</td><td>15</td></tr><tr><td>2</td><td>0.21097</td><td>0.029106</td><td>0.522927</td><td>0.317476</td><td>9.596375</td><td>14</td></tr><tr><td>3</td><td>0.701792</td><td>0.527346</td><td>0.352297</td><td>0.912383</td><td>4.874474</td><td>16</td></tr><tr><td>4</td><td>0.723815</td><td>0.544753</td><td>0.311694</td><td>0.210474</td><td>5.696281</td><td>17</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 7)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬──────────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       ┆ nb_l_inf_cnt │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---          │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ u32          │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪══════════════╡\n",
       "│ 0   ┆ 0.720043 ┆ 0.762057 ┆ 0.0802   ┆ 0.853373 ┆ 1.184888 ┆ 16           │\n",
       "│ 1   ┆ 0.746859 ┆ 0.774783 ┆ 0.969885 ┆ 0.027992 ┆ 6.011372 ┆ 15           │\n",
       "│ 2   ┆ 0.21097  ┆ 0.029106 ┆ 0.522927 ┆ 0.317476 ┆ 9.596375 ┆ 14           │\n",
       "│ 3   ┆ 0.701792 ┆ 0.527346 ┆ 0.352297 ┆ 0.912383 ┆ 4.874474 ┆ 16           │\n",
       "│ 4   ┆ 0.723815 ┆ 0.544753 ┆ 0.311694 ┆ 0.210474 ┆ 5.696281 ┆ 17           │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────┘"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get neighbor count. The point itself is always considered a neighbor to itself.\n",
    "df.with_columns(\n",
    "    pds.query_nb_cnt(\n",
    "        pl.col(\"var1\"), \"var2\", \"var3\", # Columns used as the coordinates in n-d space, str | pl.Expr \n",
    "        r = 0.1, # radius \n",
    "        dist = \"inf\", # L Infinity distance \n",
    "        parallel = True \n",
    "    ).alias(\"nb_l_inf_cnt\")\n",
    ").head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "69ad83d1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 7)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th><th>nb_l1_r_cnt</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>u32</td></tr></thead><tbody><tr><td>0</td><td>0.720043</td><td>0.762057</td><td>0.0802</td><td>0.853373</td><td>1.184888</td><td>690</td></tr><tr><td>1</td><td>0.746859</td><td>0.774783</td><td>0.969885</td><td>0.027992</td><td>6.011372</td><td>1</td></tr><tr><td>2</td><td>0.21097</td><td>0.029106</td><td>0.522927</td><td>0.317476</td><td>9.596375</td><td>56</td></tr><tr><td>3</td><td>0.701792</td><td>0.527346</td><td>0.352297</td><td>0.912383</td><td>4.874474</td><td>1289</td></tr><tr><td>4</td><td>0.723815</td><td>0.544753</td><td>0.311694</td><td>0.210474</td><td>5.696281</td><td>23</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 7)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬─────────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       ┆ nb_l1_r_cnt │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---         │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ u32         │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪═════════════╡\n",
       "│ 0   ┆ 0.720043 ┆ 0.762057 ┆ 0.0802   ┆ 0.853373 ┆ 1.184888 ┆ 690         │\n",
       "│ 1   ┆ 0.746859 ┆ 0.774783 ┆ 0.969885 ┆ 0.027992 ┆ 6.011372 ┆ 1           │\n",
       "│ 2   ┆ 0.21097  ┆ 0.029106 ┆ 0.522927 ┆ 0.317476 ┆ 9.596375 ┆ 56          │\n",
       "│ 3   ┆ 0.701792 ┆ 0.527346 ┆ 0.352297 ┆ 0.912383 ┆ 4.874474 ┆ 1289        │\n",
       "│ 4   ┆ 0.723815 ┆ 0.544753 ┆ 0.311694 ┆ 0.210474 ┆ 5.696281 ┆ 23          │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴─────────────┘"
      ]
     },
     "execution_count": 45,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.with_columns(\n",
    "    pds.query_nb_cnt(\n",
    "        \"var1\", \"var2\", \"var3\", # Columns used as the coordinates in n-d space, str | pl.Expr \n",
    "        r = pl.col(\"r\"), # radius be an expression too\n",
    "        dist = \"l1\", # L 1 distance \n",
    "        parallel = True \n",
    "    ).alias(\"nb_l1_r_cnt\")\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "ce1a2c7a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 7)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th><th>best friends</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>list[u32]</td></tr></thead><tbody><tr><td>0</td><td>0.720043</td><td>0.762057</td><td>0.0802</td><td>0.853373</td><td>1.184888</td><td>[0, 1171, … 1754]</td></tr><tr><td>1</td><td>0.746859</td><td>0.774783</td><td>0.969885</td><td>0.027992</td><td>6.011372</td><td>[1, 906, … 1751]</td></tr><tr><td>2</td><td>0.21097</td><td>0.029106</td><td>0.522927</td><td>0.317476</td><td>9.596375</td><td>[2, 50, … 853]</td></tr><tr><td>3</td><td>0.701792</td><td>0.527346</td><td>0.352297</td><td>0.912383</td><td>4.874474</td><td>[3, 1558, … 921]</td></tr><tr><td>4</td><td>0.723815</td><td>0.544753</td><td>0.311694</td><td>0.210474</td><td>5.696281</td><td>[4, 3, … 485]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 7)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬───────────────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       ┆ best friends      │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---               │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ list[u32]         │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪═══════════════════╡\n",
       "│ 0   ┆ 0.720043 ┆ 0.762057 ┆ 0.0802   ┆ 0.853373 ┆ 1.184888 ┆ [0, 1171, … 1754] │\n",
       "│ 1   ┆ 0.746859 ┆ 0.774783 ┆ 0.969885 ┆ 0.027992 ┆ 6.011372 ┆ [1, 906, … 1751]  │\n",
       "│ 2   ┆ 0.21097  ┆ 0.029106 ┆ 0.522927 ┆ 0.317476 ┆ 9.596375 ┆ [2, 50, … 853]    │\n",
       "│ 3   ┆ 0.701792 ┆ 0.527346 ┆ 0.352297 ┆ 0.912383 ┆ 4.874474 ┆ [3, 1558, … 921]  │\n",
       "│ 4   ┆ 0.723815 ┆ 0.544753 ┆ 0.311694 ┆ 0.210474 ┆ 5.696281 ┆ [4, 3, … 485]     │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴───────────────────┘"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get ids of the k nearest neighbors. \n",
    "# The point itself is always considered a neighbor to itself, so k + 1 elements will be returned.\n",
    "df.with_columns(\n",
    "    pds.query_knn_ptwise(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), pl.col(\"var3\"), # Columns used as the coordinates in n-d space\n",
    "        index = \"id\",  # pl.col(\"id\"), str | pl.Expr\n",
    "        k = 3, \n",
    "        dist = \"l2\", # squared l2\n",
    "        parallel = True\n",
    "    ).alias(\"best friends\")\n",
    ").head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "67a769f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (5, 3)\n",
      "┌─────┬──────────────────┬────────────────────┐\n",
      "│ id  ┆ best friends     ┆ best friends count │\n",
      "│ --- ┆ ---              ┆ ---                │\n",
      "│ u32 ┆ list[u32]        ┆ u32                │\n",
      "╞═════╪══════════════════╪════════════════════╡\n",
      "│ 0   ┆ [0, 1171, … 912] ┆ 9                  │\n",
      "│ 1   ┆ [1, 906, … 831]  ┆ 5                  │\n",
      "│ 2   ┆ [2, 50, … 1682]  ┆ 8                  │\n",
      "│ 3   ┆ [3, 1558, … 66]  ┆ 7                  │\n",
      "│ 4   ┆ [4, 3, … 1370]   ┆ 6                  │\n",
      "└─────┴──────────────────┴────────────────────┘\n"
     ]
    }
   ],
   "source": [
    "# Get all neighbors within radius r, call them best friends\n",
    "print(\n",
    "\n",
    "df.select(\n",
    "    pl.col(\"id\"),\n",
    "    pds.query_radius_ptwise(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), pl.col(\"var3\"), # Columns used as the coordinates in 3d space\n",
    "        index = pl.col(\"id\"),\n",
    "        r = 0.1, \n",
    "        dist = \"l2\", # actually this is squared l2\n",
    "        parallel = True\n",
    "    ).alias(\"best friends\"),\n",
    ").with_columns( # -1 to remove the point itself\n",
    "    (pl.col(\"best friends\").list.len() - 1).alias(\"best friends count\")\n",
    ").head()\n",
    "\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "06cd7fb2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 8)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th><th>idx</th><th>dist</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>list[u32]</td><td>list[f64]</td></tr></thead><tbody><tr><td>0</td><td>0.720043</td><td>0.762057</td><td>0.0802</td><td>0.853373</td><td>1.184888</td><td>[0, 1171, … 1754]</td><td>[0.0, 0.054683, … 0.077248]</td></tr><tr><td>1</td><td>0.746859</td><td>0.774783</td><td>0.969885</td><td>0.027992</td><td>6.011372</td><td>[1, 906, … 1751]</td><td>[0.0, 0.042337, … 0.053288]</td></tr><tr><td>2</td><td>0.21097</td><td>0.029106</td><td>0.522927</td><td>0.317476</td><td>9.596375</td><td>[2, 50, … 853]</td><td>[0.0, 0.059335, … 0.06505]</td></tr><tr><td>3</td><td>0.701792</td><td>0.527346</td><td>0.352297</td><td>0.912383</td><td>4.874474</td><td>[3, 1558, … 921]</td><td>[0.0, 0.015422, … 0.067852]</td></tr><tr><td>4</td><td>0.723815</td><td>0.544753</td><td>0.311694</td><td>0.210474</td><td>5.696281</td><td>[4, 3, … 485]</td><td>[0.0, 0.049363, … 0.060237]</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 8)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┬──────────────────┬──────────────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       ┆ idx              ┆ dist             │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---              ┆ ---              │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ list[u32]        ┆ list[f64]        │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╪══════════════════╪══════════════════╡\n",
       "│ 0   ┆ 0.720043 ┆ 0.762057 ┆ 0.0802   ┆ 0.853373 ┆ 1.184888 ┆ [0, 1171, …      ┆ [0.0, 0.054683,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆ 1754]            ┆ … 0.077248]      │\n",
       "│ 1   ┆ 0.746859 ┆ 0.774783 ┆ 0.969885 ┆ 0.027992 ┆ 6.011372 ┆ [1, 906, … 1751] ┆ [0.0, 0.042337,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.053288]      │\n",
       "│ 2   ┆ 0.21097  ┆ 0.029106 ┆ 0.522927 ┆ 0.317476 ┆ 9.596375 ┆ [2, 50, … 853]   ┆ [0.0, 0.059335,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.06505]       │\n",
       "│ 3   ┆ 0.701792 ┆ 0.527346 ┆ 0.352297 ┆ 0.912383 ┆ 4.874474 ┆ [3, 1558, … 921] ┆ [0.0, 0.015422,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.067852]      │\n",
       "│ 4   ┆ 0.723815 ┆ 0.544753 ┆ 0.311694 ┆ 0.210474 ┆ 5.696281 ┆ [4, 3, … 485]    ┆ [0.0, 0.049363,  │\n",
       "│     ┆          ┆          ┆          ┆          ┆          ┆                  ┆ … 0.060237]      │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┴──────────────────┴──────────────────┘"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Get ids of the k nearest neighbors and distances\n",
    "# The point itself is always considered a neighbor to itself, so k + 1 elements will be returned.\n",
    "df.with_columns(\n",
    "    pds.query_knn_ptwise(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), pl.col(\"var3\"), # Columns used as the coordinates in n-d space\n",
    "        index = pl.col(\"id\"),\n",
    "        k = 3, \n",
    "        dist = \"l2\", # actually this is squared l2\n",
    "        parallel = True,\n",
    "        return_dist = True\n",
    "    ).alias(\"best_friends_w_dist\")\n",
    ").unnest(\"best_friends_w_dist\").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "b5c69ae4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>3</td><td>0.701792</td><td>0.527346</td><td>0.352297</td><td>0.912383</td><td>4.874474</td></tr><tr><td>4</td><td>0.723815</td><td>0.544753</td><td>0.311694</td><td>0.210474</td><td>5.696281</td></tr><tr><td>6</td><td>0.736724</td><td>0.776174</td><td>0.693574</td><td>0.532166</td><td>3.944928</td></tr><tr><td>7</td><td>0.617642</td><td>0.788939</td><td>0.488318</td><td>0.27519</td><td>2.900536</td></tr><tr><td>8</td><td>0.327207</td><td>0.456177</td><td>0.469014</td><td>0.128877</td><td>9.917232</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 3   ┆ 0.701792 ┆ 0.527346 ┆ 0.352297 ┆ 0.912383 ┆ 4.874474 │\n",
       "│ 4   ┆ 0.723815 ┆ 0.544753 ┆ 0.311694 ┆ 0.210474 ┆ 5.696281 │\n",
       "│ 6   ┆ 0.736724 ┆ 0.776174 ┆ 0.693574 ┆ 0.532166 ┆ 3.944928 │\n",
       "│ 7   ┆ 0.617642 ┆ 0.788939 ┆ 0.488318 ┆ 0.27519  ┆ 2.900536 │\n",
       "│ 8   ┆ 0.327207 ┆ 0.456177 ┆ 0.469014 ┆ 0.128877 ┆ 9.917232 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Filter to only points near the given point\n",
    "df.filter(\n",
    "    pds.within_dist_from(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), pl.col(\"var3\"), # Columns used as the coordinates in n-d space\n",
    "        pt = [0.5, 0.5, 0.5],\n",
    "        r = 0.2,\n",
    "        dist = \"l2\" # actually this is squared l2, so this is asking for squared l2 <= 0.2\n",
    "    )\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "4ab9e8f8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>36</td><td>0.568433</td><td>0.456969</td><td>0.05886</td><td>0.006766</td><td>0.753822</td></tr><tr><td>117</td><td>0.521121</td><td>0.521892</td><td>0.129322</td><td>0.377708</td><td>0.731702</td></tr><tr><td>138</td><td>0.531877</td><td>0.488015</td><td>0.040861</td><td>0.905918</td><td>1.466601</td></tr><tr><td>169</td><td>0.465704</td><td>0.501999</td><td>0.253489</td><td>0.625459</td><td>6.367421</td></tr><tr><td>176</td><td>0.473739</td><td>0.491308</td><td>0.894173</td><td>0.652556</td><td>6.084469</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 36  ┆ 0.568433 ┆ 0.456969 ┆ 0.05886  ┆ 0.006766 ┆ 0.753822 │\n",
       "│ 117 ┆ 0.521121 ┆ 0.521892 ┆ 0.129322 ┆ 0.377708 ┆ 0.731702 │\n",
       "│ 138 ┆ 0.531877 ┆ 0.488015 ┆ 0.040861 ┆ 0.905918 ┆ 1.466601 │\n",
       "│ 169 ┆ 0.465704 ┆ 0.501999 ┆ 0.253489 ┆ 0.625459 ┆ 6.367421 │\n",
       "│ 176 ┆ 0.473739 ┆ 0.491308 ┆ 0.894173 ┆ 0.652556 ┆ 6.084469 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Haversine distance is available when dimension is 2\n",
    "df.filter(\n",
    "    pds.within_dist_from(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), # Columns used as the coordinates in n-d space\n",
    "        pt = [0.5, 0.5],\n",
    "        r = 10, # in km\n",
    "        dist = \"h\" \n",
    "    )\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "7d3f5ae4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>var1</th><th>var2</th><th>var3</th><th>r</th><th>rh</th></tr><tr><td>u32</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>169</td><td>0.465704</td><td>0.501999</td><td>0.253489</td><td>0.625459</td><td>6.367421</td></tr><tr><td>176</td><td>0.473739</td><td>0.491308</td><td>0.894173</td><td>0.652556</td><td>6.084469</td></tr><tr><td>235</td><td>0.564589</td><td>0.484392</td><td>0.057102</td><td>0.975023</td><td>8.699902</td></tr><tr><td>367</td><td>0.478239</td><td>0.566379</td><td>0.620646</td><td>0.384922</td><td>9.836408</td></tr><tr><td>383</td><td>0.501891</td><td>0.46347</td><td>0.135889</td><td>0.616873</td><td>7.838947</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌─────┬──────────┬──────────┬──────────┬──────────┬──────────┐\n",
       "│ id  ┆ var1     ┆ var2     ┆ var3     ┆ r        ┆ rh       │\n",
       "│ --- ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │\n",
       "│ u32 ┆ f64      ┆ f64      ┆ f64      ┆ f64      ┆ f64      │\n",
       "╞═════╪══════════╪══════════╪══════════╪══════════╪══════════╡\n",
       "│ 169 ┆ 0.465704 ┆ 0.501999 ┆ 0.253489 ┆ 0.625459 ┆ 6.367421 │\n",
       "│ 176 ┆ 0.473739 ┆ 0.491308 ┆ 0.894173 ┆ 0.652556 ┆ 6.084469 │\n",
       "│ 235 ┆ 0.564589 ┆ 0.484392 ┆ 0.057102 ┆ 0.975023 ┆ 8.699902 │\n",
       "│ 367 ┆ 0.478239 ┆ 0.566379 ┆ 0.620646 ┆ 0.384922 ┆ 9.836408 │\n",
       "│ 383 ┆ 0.501891 ┆ 0.46347  ┆ 0.135889 ┆ 0.616873 ┆ 7.838947 │\n",
       "└─────┴──────────┴──────────┴──────────┴──────────┴──────────┘"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.filter(\n",
    "    pds.within_dist_from(\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), \n",
    "        pt = [0.5, 0.5],\n",
    "        # radius can also be an existing column in the dataframe.\n",
    "        r = pl.col(\"rh\"), \n",
    "        dist = \"h\" \n",
    "    )\n",
    ").head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "f14627bc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 3)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>id</th><th>friends</th><th>count</th></tr><tr><td>u64</td><td>list[u32]</td><td>u32</td></tr></thead><tbody><tr><td>0</td><td>[0, 1345, … 304]</td><td>4</td></tr><tr><td>1</td><td>[1, 6, 278]</td><td>3</td></tr><tr><td>2</td><td>[2, 934, 853]</td><td>3</td></tr><tr><td>3</td><td>[3, 1584, … 159]</td><td>5</td></tr><tr><td>4</td><td>[4, 1939, … 392]</td><td>5</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 3)\n",
       "┌─────┬──────────────────┬───────┐\n",
       "│ id  ┆ friends          ┆ count │\n",
       "│ --- ┆ ---              ┆ ---   │\n",
       "│ u64 ┆ list[u32]        ┆ u32   │\n",
       "╞═════╪══════════════════╪═══════╡\n",
       "│ 0   ┆ [0, 1345, … 304] ┆ 4     │\n",
       "│ 1   ┆ [1, 6, 278]      ┆ 3     │\n",
       "│ 2   ┆ [2, 934, 853]    ┆ 3     │\n",
       "│ 3   ┆ [3, 1584, … 159] ┆ 5     │\n",
       "│ 4   ┆ [4, 1939, … 392] ┆ 5     │\n",
       "└─────┴──────────────────┴───────┘"
      ]
     },
     "execution_count": 52,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "friends = df.select(\n",
    "    pl.col(\"id\").cast(pl.UInt64),\n",
    "    pds.query_radius_ptwise(\n",
    "        # Columns used as the coordinates in n-d space\n",
    "        pl.col(\"var1\"), pl.col(\"var2\"), \n",
    "        index=pl.col(\"id\"),\n",
    "        r = 0.02, \n",
    "        dist = \"l2\",\n",
    "    ).alias(\"friends\")\n",
    ").with_columns(\n",
    "    pl.col(\"friends\").list.len().alias(\"count\")\n",
    ")\n",
    "friends.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98264a05",
   "metadata": {},
   "source": [
    "# Using PDS Expressions On pl.Series, NumPy arrays, or pd.Series, etc.\n",
    "\n",
    "The output by default is always a Polars Series. The user gets to choose whether to turn it into NumPy, Pandas, or other data structures."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "0a42a771",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_19683/3354819425.py:3: UserWarning: The compatibility layer is considered experimental.\n",
      "  from polars_ds.compat import compat as pds2\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 6)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>actual</th><th>predicted</th><th>0-2</th><th>0-9</th><th>s1</th><th>s2</th></tr><tr><td>f64</td><td>f64</td><td>i32</td><td>i32</td><td>str</td><td>str</td></tr></thead><tbody><tr><td>1.0</td><td>0.78908</td><td>0</td><td>1</td><td>&quot;7J&quot;</td><td>&quot;k&quot;</td></tr><tr><td>1.0</td><td>0.503485</td><td>1</td><td>5</td><td>&quot;S&quot;</td><td>&quot;yj&quot;</td></tr><tr><td>1.0</td><td>0.736868</td><td>2</td><td>8</td><td>&quot;iB&quot;</td><td>&quot;p&quot;</td></tr><tr><td>1.0</td><td>0.904397</td><td>1</td><td>4</td><td>&quot;R&quot;</td><td>&quot;Js&quot;</td></tr><tr><td>1.0</td><td>0.843379</td><td>1</td><td>2</td><td>&quot;A&quot;</td><td>&quot;WR&quot;</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 6)\n",
       "┌────────┬───────────┬─────┬─────┬─────┬─────┐\n",
       "│ actual ┆ predicted ┆ 0-2 ┆ 0-9 ┆ s1  ┆ s2  │\n",
       "│ ---    ┆ ---       ┆ --- ┆ --- ┆ --- ┆ --- │\n",
       "│ f64    ┆ f64       ┆ i32 ┆ i32 ┆ str ┆ str │\n",
       "╞════════╪═══════════╪═════╪═════╪═════╪═════╡\n",
       "│ 1.0    ┆ 0.78908   ┆ 0   ┆ 1   ┆ 7J  ┆ k   │\n",
       "│ 1.0    ┆ 0.503485  ┆ 1   ┆ 5   ┆ S   ┆ yj  │\n",
       "│ 1.0    ┆ 0.736868  ┆ 2   ┆ 8   ┆ iB  ┆ p   │\n",
       "│ 1.0    ┆ 0.904397  ┆ 1   ┆ 4   ┆ R   ┆ Js  │\n",
       "│ 1.0    ┆ 0.843379  ┆ 1   ┆ 2   ┆ A   ┆ WR  │\n",
       "└────────┴───────────┴─────┴─────┴─────┴─────┘"
      ]
     },
     "execution_count": 53,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from polars_ds.compat import compat as pds2\n",
    "\n",
    "df = pds.frame(size=100_000).select(\n",
    "    pds.random(0.0, 1.0).round().alias(\"actual\"),\n",
    "    pds.random(0.0, 1.0).alias(\"predicted\"),\n",
    "    pds.random_int(0, 3).alias(\"0-2\"),\n",
    "    pds.random_int(0, 10).alias(\"0-9\"),\n",
    "    pds.random_str(min_size=1, max_size=2).alias(\"s1\"),\n",
    "    pds.random_str(min_size=1, max_size=2).alias(\"s2\"),\n",
    ")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "e6574a57",
   "metadata": {},
   "outputs": [],
   "source": [
    "df_pd = df.to_pandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "70830d04",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1,)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>jaccard_col</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.3</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1,)\n",
       "Series: 'jaccard_col' [f64]\n",
       "[\n",
       "\t0.3\n",
       "]"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Pandas Series\n",
    "pds2.jaccard_col(df_pd[\"0-2\"], df_pd[\"0-9\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "471cbbbd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "shape: (1, 5)\n",
      "┌───────────┬──────────┬──────────┬───────────────────┬──────────┐\n",
      "│ precision ┆ recall   ┆ f        ┆ average_precision ┆ roc_auc  │\n",
      "│ ---       ┆ ---      ┆ ---      ┆ ---               ┆ ---      │\n",
      "│ f64       ┆ f64      ┆ f64      ┆ f64               ┆ f64      │\n",
      "╞═══════════╪══════════╪══════════╪═══════════════════╪══════════╡\n",
      "│ 0.503697  ┆ 0.500827 ┆ 0.502258 ┆ 0.50319           ┆ 0.501299 │\n",
      "└───────────┴──────────┴──────────┴───────────────────┴──────────┘\n",
      "shape: (1, 5)\n",
      "┌───────────┬──────────┬──────────┬───────────────────┬──────────┐\n",
      "│ precision ┆ recall   ┆ f        ┆ average_precision ┆ roc_auc  │\n",
      "│ ---       ┆ ---      ┆ ---      ┆ ---               ┆ ---      │\n",
      "│ f64       ┆ f64      ┆ f64      ┆ f64               ┆ f64      │\n",
      "╞═══════════╪══════════╪══════════╪═══════════════════╪══════════╡\n",
      "│ 0.503697  ┆ 0.500827 ┆ 0.502258 ┆ 0.50319           ┆ 0.501299 │\n",
      "└───────────┴──────────┴──────────┴───────────────────┴──────────┘\n",
      "shape: (1, 5)\n",
      "┌───────────┬──────────┬──────────┬───────────────────┬──────────┐\n",
      "│ precision ┆ recall   ┆ f        ┆ average_precision ┆ roc_auc  │\n",
      "│ ---       ┆ ---      ┆ ---      ┆ ---               ┆ ---      │\n",
      "│ f64       ┆ f64      ┆ f64      ┆ f64               ┆ f64      │\n",
      "╞═══════════╪══════════╪══════════╪═══════════════════╪══════════╡\n",
      "│ 0.503697  ┆ 0.500827 ┆ 0.502258 ┆ 0.50319           ┆ 0.501299 │\n",
      "└───────────┴──────────┴──────────┴───────────────────┴──────────┘\n"
     ]
    }
   ],
   "source": [
    "# Polars Series\n",
    "print(pds2.query_binary_metrics(df[\"actual\"], df[\"predicted\"]).struct.unnest())\n",
    "# Pyarrow Series\n",
    "print(pds2.query_binary_metrics(df[\"actual\"].to_arrow(), df[\"predicted\"].to_arrow()).struct.unnest())\n",
    "# NumPy\n",
    "print(pds2.query_binary_metrics(df[\"actual\"].to_numpy(), df[\"predicted\"].to_numpy()).struct.unnest())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "9d15fc63",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (5, 4)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>&lt;=</th><th>baseline_pct</th><th>actual_pct</th><th>psi_bin</th></tr><tr><td>f64</td><td>f64</td><td>f64</td><td>f64</td></tr></thead><tbody><tr><td>0.194464</td><td>0.2</td><td>0.187</td><td>0.000874</td></tr><tr><td>0.413897</td><td>0.2</td><td>0.207</td><td>0.000241</td></tr><tr><td>0.5999</td><td>0.2</td><td>0.19</td><td>0.000513</td></tr><tr><td>0.804026</td><td>0.2</td><td>0.202</td><td>0.00002</td></tr><tr><td>inf</td><td>0.2</td><td>0.214</td><td>0.000947</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (5, 4)\n",
       "┌──────────┬──────────────┬────────────┬──────────┐\n",
       "│ <=       ┆ baseline_pct ┆ actual_pct ┆ psi_bin  │\n",
       "│ ---      ┆ ---          ┆ ---        ┆ ---      │\n",
       "│ f64      ┆ f64          ┆ f64        ┆ f64      │\n",
       "╞══════════╪══════════════╪════════════╪══════════╡\n",
       "│ 0.194464 ┆ 0.2          ┆ 0.187      ┆ 0.000874 │\n",
       "│ 0.413897 ┆ 0.2          ┆ 0.207      ┆ 0.000241 │\n",
       "│ 0.5999   ┆ 0.2          ┆ 0.19       ┆ 0.000513 │\n",
       "│ 0.804026 ┆ 0.2          ┆ 0.202      ┆ 0.00002  │\n",
       "│ inf      ┆ 0.2          ┆ 0.214      ┆ 0.000947 │\n",
       "└──────────┴──────────────┴────────────┴──────────┘"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# NumPy Arrays\n",
    "pds2.psi(\n",
    "    np.random.random(size = 1000), \n",
    "    np.random.random(size = 1000), \n",
    "    n_bins = 5, \n",
    "    return_report = True\n",
    ").struct.unnest()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "8c5c3fee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1,)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>cid_ce</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>12.900387</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1,)\n",
       "Series: 'cid_ce' [f64]\n",
       "[\n",
       "\t12.900387\n",
       "]"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds2.query_cid_ce(\n",
    "    np.random.random(size = 1000),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "b0c14067",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (1,)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>c3_stats</th></tr><tr><td>f64</td></tr></thead><tbody><tr><td>0.123215</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (1,)\n",
       "Series: 'c3_stats' [f64]\n",
       "[\n",
       "\t0.123215\n",
       "]"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds2.query_c3_stats(\n",
    "    pl.Series(values=np.random.random(size = 1000)), \n",
    "    lag = 3\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "f7cc8f20",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><style>\n",
       ".dataframe > thead > tr,\n",
       ".dataframe > tbody > tr {\n",
       "  text-align: right;\n",
       "  white-space: pre-wrap;\n",
       "}\n",
       "</style>\n",
       "<small>shape: (10,)</small><table border=\"1\" class=\"dataframe\"><thead><tr><th>str_leven</th></tr><tr><td>u32</td></tr></thead><tbody><tr><td>2</td></tr><tr><td>2</td></tr><tr><td>2</td></tr><tr><td>2</td></tr><tr><td>2</td></tr><tr><td>2</td></tr><tr><td>1</td></tr><tr><td>2</td></tr><tr><td>2</td></tr><tr><td>2</td></tr></tbody></table></div>"
      ],
      "text/plain": [
       "shape: (10,)\n",
       "Series: 'str_leven' [u32]\n",
       "[\n",
       "\t2\n",
       "\t2\n",
       "\t2\n",
       "\t2\n",
       "\t2\n",
       "\t2\n",
       "\t1\n",
       "\t2\n",
       "\t2\n",
       "\t2\n",
       "]"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pds2.str_leven(df[\"s1\"], df[\"s2\"]).head() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2aab4a3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
